from scipy import constants, integrate
import numpy as np
from numpy.random.mtrand import RandomState 

proton_mass = constants.proton_mass/ constants.electron_volt * constants.speed_of_light**2 
speed_of_light = constants.speed_of_light

def uniform(low,high,size,seed=None):
    A = RandomState(seed)
    return A.uniform(low,high,size)

def normal(mean, sig, size, seed=None, cutoff=None):
    A = RandomState(seed)
    g = A.normal(mean, sig, size)
    if cutoff:
        ind=np.flatnonzero((np.abs(g-mean)) > (cutoff*sig))
        while len(ind)!=0:
            g[ind] = A.normal(mean,sig,len(ind))
            ind=np.flatnonzero((np.abs(g-mean)) > (cutoff*sig))
    return g

def IntET(beta,freq,Ez,z):
    k = freq*2*np.pi/(beta*speed_of_light)
    V = integrate.trapz(Ez*np.cos(z*k),dx=z[1]-z[0])
    return V

def WToBeta(w):
    gamma = 1 + w/proton_mass
    beta = np.sqrt(1 - 1/gamma**2)
    return beta
                
def BetaToW(beta):
    gamma = np.sqrt(1/(1-beta**2))
    w = proton_mass*(gamma-1)
    return w

def loadLinacInfo(file,start,end):
    '''Extract drift lengths from TraceWin file'''
    Drifts = []
    f = open(file)
    lines=f.readlines()
    f.close()
    start_line=[lines.index(line) for line in lines if start in line if line.split()[0]==start][0]
    end_line=[lines.index(line) for line in lines if end in line if line.split()[0]==end][0]
    lines = lines[start_line:end_line]
    TP,TLinac = [],[]
    d = 0
    for line in lines:
        if ('DRIFT' in line or 'QUAD' in line) & (line[0]!=';'):
            d += float(line.split()[3])
            if 'QUAD' in line:
                TP.append(quadMat(float(line.split()[4]),float(line.split()[3])*1e-3))                
            elif 'DRIFT' in line:
                TP.append(driftMat(float(line.split()[3])*1e-3))    
        if 'NCELLS' in line:
            Drifts.append(d)
            d = 0
            TLinac.append(MatrixDot(TP))
            TP = []
    Drifts.append(d)
    return np.array(Drifts)*1e-3,np.array(TLinac)

def meminfo():
    VmRSS = [float(i.split()[1]) for i in open('/proc/self/status') if 'VmRSS' in i][0]
    print 'Memory used by process', VmRSS/1024.0, 'MB'

def driftMat(l):
    d = np.ones((2,2))
    d[1,0]=0
    d[0,1]=l
    return d

def quadMat(G,l):
    P=1
    T = np.ones((2,2))
    if G > 0:
        K = np.sqrt(G/P)
        C = np.cos(K*l)
        S = np.sin(K*l)
        T[0,0] = C
        T[1,0] = -K*S
        T[0,1] = S/K
        T[1,1] = C
    elif G<0:
        K = np.sqrt(-G/P)
        C = np.cosh(K*l)
        S = np.sinh(K*l)
        T[0,0] = C
        T[1,0] = K*S
        T[0,1] = S/K
        T[1,1] = C
    return T

def MatrixDot(lst):
    lst.reverse()
    stack=1
    for i in xrange(len(lst)):
        answ = np.dot(lst[i],stack)
        stack = 1*answ
    return stack

def MatrixDotArray(arr):
    arr=arr[::-1]
    stack=1
    for i in xrange(arr.shape[0]):
        answ = np.dot(arr[i],stack)
        stack = 1*answ
    return stack

def extractTMatrix(file):
    eleNum=len([line for line in open(file) if '#' in line])
    data = np.array([line.split() for line in open(file) if 'm' not in line],dtype=float).reshape(eleNum,6,6)
    return data

def Interpolate(x,y,period=None):
    '''Perform interpolation''' 
    if period!=None:
        period = np.array(period,dtype='int')
    else:
        period = np.ones(y.size,dtype='int')*y.size
    Corr = np.zeros(y.size)
    i=0
    while i < y.size:
        Corr[i] = (y[i+1] - y[i])/(x[i+1]-x[i]) 
        i+=1
        for j in xrange(1,period[i]-1): 
            Corr[i] = (0.5*((y[i]-y[i-1])/(x[i] - x[i-1])) + 0.5*((y[i+1]-y[i])/(x[i+1]-x[i])))
            i+=1
        Corr[i] = (y[i] - y[i-1])/(x[i] - x[i-1])
        i+=1
    return Corr

